import json
import numpy as np
from sklearn.model_selection import train_test_split
import tensorflow.keras as keras
import matplotlib.pyplot as plt


DATASET_PATH = "D:\\Audio\\datasets\\genres\\data_10.json"


def load_data(dataset_path):
    with open(dataset_path, "r") as fp:
        data = json.load(fp)
        
    # convert lists into numpy arrays
    X = np.array(data["mfcc"])
    y = np.array(data["labels"])
    
    return X, y 


def plot_history(history):
    
    fig, axs = plt.subplots(2)
    
    # create accuracy subplot
    axs[0].plot(history.history["acc"], label="train accuracy")
    axs[0].plot(history.history["val_acc"], label="test accuracy")
    axs[0].set_ylabel("Accuracy")
    axs[0].set_xlabel("Epoch")
    axs[0].legend(loc="lower right")
    axs[0].set_title("Accuracy eval")
    
    # create error subplot
    axs[1].plot(history.history["loss"], label="train error")
    axs[1].plot(history.history["val_loss"], label="test error")
    axs[1].set_ylabel("Error")
    axs[1].set_xlabel("Epoch")
    axs[1].legend(loc="upper right")
    axs[1].set_title("Error eval")
    
    plt.show()
    
    
def prepare_datasets(test_size, validation_size):
    
    # load data
    X, y = load_data(DATASET_PATH)
    
    # create train/test split 
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size)
    
    # create train/validation split
    X_train, X_validation, y_train, y_validation = train_test_split(X_train, y_train, test_size=validation_size)
    
    # dont need new axis
    # 3d array -> (130, 13, 1)
    # X_train = X_train[..., np.newaxis]  # 4d array -> (num_samples, 130, 13, 1)
    # X_test = X_test[..., np.newaxis]
    # X_validation = X_validation[..., np.newaxis]
    
    return X_train, X_validation, X_test, y_train, y_validation, y_test
    
    
def build_model(input_shape):
    """
    RNN-LSTM model
    """
    
    # create model
    model = keras.Sequential()
    
    # 2 LSTM layers
    # CuDNNLSTM in CUDA
    # model.add(keras.layers.CuDNNLSTM(64, input_shape=input_shape, return_sequences=True))
    # model.add(keras.layers.CuDNNLSTM(64))
    model.add(keras.layers.LSTM(64, input_shape=input_shape, return_sequences=True))
    model.add(keras.layers.LSTM(64))
    
    # dense layer
    model.add(keras.layers.Dense(64, activation='relu'))
    model.add(keras.layers.Dropout(0.3))
    
    # output layer
    model.add(keras.layers.Dense(10, activation='softmax'))  
    return model
        

def predict(model, X, y):
    
    X = X[np.newaxis, ...]
    
    # prediction = [ [0.1, 0.2, ...]]
    prediction = model.predict(X)  # X -> (1, 130, 13, 1)
    
    # extract index with max value
    predicted_index = np.argmax(prediction, axis=1)
    
    print("Expected index: {}, Predicted index: {}".format(y, predicted_index))
           
           
if __name__ == "__main__":
    
    # create train, validation and test sets
    X_train, X_validation, X_test, y_train, y_validation, y_test = prepare_datasets(0.25, 0.2)
    
    # build the RNN-LSTM network
    input_shape = (X_train.shape[1], X_train.shape[2])  # 130, 13
    model = build_model(input_shape)
    
    # compile the CNN 
    optimizer = keras.optimizers.Adam(lr=0.0001)
    model.compile(optimizer=optimizer,
                  loss="sparse_categorical_crossentropy",
                  metrics=["accuracy"])
    model.summary()
    
    # train the CNN
    history = model.fit(X_train, y_train, 
              validation_data=(X_validation, y_validation), 
              batch_size=32, 
              epochs=30)
    

    
    # evalutate the CNN on test set 
    test_loss, test_acc = model.evaluate(X_test, y_test, verbose=1)
    print("Accuracy on test set is {}".format(test_acc))
    
    # make prediction on a sample
    X = X_test[100]
    y = y_test[100]
    predict(model, X, y)
    
    plot_history(history)
    